package com.sjw.workflow.dag;

import com.sjw.workflow.core.WorkflowContext;
import java.util.Map;
import java.util.HashMap;
import java.util.List;
import java.util.ArrayList;
import java.util.Set;
import java.util.HashSet;
import java.util.Queue;
import java.util.LinkedList;
import java.util.stream.Collectors;

/**
 * Workflow DAG图，用于管理Agent的依赖关系和执行顺序
 */
public class WorkflowDAG {
    
    private final Map<String, WorkflowNode> nodes;
    private final Map<String, List<String>> adjacencyList;
    private final String workflowId;
    
    public WorkflowDAG(String workflowId) {
        this.workflowId = workflowId;
        this.nodes = new HashMap<>();
        this.adjacencyList = new HashMap<>();
    }
    
    /**
     * 添加节点到DAG
     */
    public void addNode(WorkflowNode node) {
        nodes.put(node.getId(), node);
        adjacencyList.put(node.getId(), new ArrayList<>());
    }
    
    /**
     * 添加边（依赖关系）
     */
    public void addEdge(String fromNodeId, String toNodeId) {
        if (!nodes.containsKey(fromNodeId) || !nodes.containsKey(toNodeId)) {
            throw new IllegalArgumentException("Node not found: " + fromNodeId + " or " + toNodeId);
        }
        
        // 检查是否会形成环
        if (wouldCreateCycle(fromNodeId, toNodeId)) {
            throw new IllegalArgumentException("Adding edge would create a cycle: " + fromNodeId + " -> " + toNodeId);
        }
        
        adjacencyList.get(toNodeId).add(fromNodeId);
        nodes.get(toNodeId).addDependency(fromNodeId);
    }
    
    /**
     * 检查添加边是否会形成环
     */
    private boolean wouldCreateCycle(String fromNodeId, String toNodeId) {
        // 使用DFS检查从toNodeId是否能到达fromNodeId
        Set<String> visited = new HashSet<>();
        return hasPath(toNodeId, fromNodeId, visited);
    }
    
    /**
     * 检查是否存在路径
     */
    private boolean hasPath(String start, String end, Set<String> visited) {
        if (start.equals(end)) {
            return true;
        }
        
        if (visited.contains(start)) {
            return false;
        }
        
        visited.add(start);
        
        List<String> neighbors = adjacencyList.get(start);
        if (neighbors != null) {
            for (String neighbor : neighbors) {
                if (hasPath(neighbor, end, visited)) {
                    return true;
                }
            }
        }
        
        return false;
    }
    
    /**
     * 获取可以执行的节点（没有依赖或依赖已满足）
     */
    public List<WorkflowNode> getReadyNodes(WorkflowContext context) {
        return nodes.values().stream()
                .filter(node -> node.canExecute(context))
                .collect(Collectors.toList());
    }
    
    /**
     * 获取所有节点
     */
    public List<WorkflowNode> getAllNodes() {
        return new ArrayList<>(nodes.values());
    }
    
    /**
     * 获取指定节点
     */
    public WorkflowNode getNode(String nodeId) {
        return nodes.get(nodeId);
    }
    
    /**
     * 检查DAG是否为空
     */
    public boolean isEmpty() {
        return nodes.isEmpty();
    }
    
    /**
     * 获取节点数量
     */
    public int size() {
        return nodes.size();
    }
    
    /**
     * 拓扑排序，返回执行顺序
     */
    public List<String> topologicalSort() {
        List<String> result = new ArrayList<>();
        Map<String, Integer> inDegree = new HashMap<>();
        Queue<String> queue = new LinkedList<>();
        
        // 计算入度
        for (String nodeId : nodes.keySet()) {
            inDegree.put(nodeId, 0);
        }
        
        for (String nodeId : nodes.keySet()) {
            List<String> dependencies = adjacencyList.get(nodeId);
            if (dependencies != null) {
                for (String dep : dependencies) {
                    inDegree.put(nodeId, inDegree.get(nodeId) + 1);
                }
            }
        }
        
        // 找到入度为0的节点
        for (Map.Entry<String, Integer> entry : inDegree.entrySet()) {
            if (entry.getValue() == 0) {
                queue.offer(entry.getKey());
            }
        }
        
        // 拓扑排序
        while (!queue.isEmpty()) {
            String current = queue.poll();
            result.add(current);
            
            // 更新相邻节点的入度
            for (String nodeId : nodes.keySet()) {
                List<String> dependencies = adjacencyList.get(nodeId);
                if (dependencies != null && dependencies.contains(current)) {
                    inDegree.put(nodeId, inDegree.get(nodeId) - 1);
                    if (inDegree.get(nodeId) == 0) {
                        queue.offer(nodeId);
                    }
                }
            }
        }
        
        return result;
    }
    
    /**
     * 验证DAG结构
     */
    public boolean validate() {
        // 检查是否有环
        Set<String> visited = new HashSet<>();
        Set<String> recursionStack = new HashSet<>();
        
        for (String nodeId : nodes.keySet()) {
            if (!visited.contains(nodeId)) {
                if (hasCycle(nodeId, visited, recursionStack)) {
                    return false;
                }
            }
        }
        
        return true;
    }
    
    /**
     * 检查是否有环
     */
    private boolean hasCycle(String nodeId, Set<String> visited, Set<String> recursionStack) {
        visited.add(nodeId);
        recursionStack.add(nodeId);
        
        List<String> dependencies = adjacencyList.get(nodeId);
        if (dependencies != null) {
            for (String dep : dependencies) {
                if (!visited.contains(dep)) {
                    if (hasCycle(dep, visited, recursionStack)) {
                        return true;
                    }
                } else if (recursionStack.contains(dep)) {
                    return true;
                }
            }
        }
        
        recursionStack.remove(nodeId);
        return false;
    }
    
    /**
     * 获取Workflow ID
     */
    public String getWorkflowId() {
        return workflowId;
    }
}
